#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author ：hhx
@Date ：2022/5/25 14:14 
@Description ：根据属性表拆分shp文件
"""

import os

import osgeo.ogr as ogr
import osgeo.osr as osr
from osgeo import gdal

# 支持中文路径
gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
# 属性表字段支持中文
gdal.SetConfigOption("SHAPE_ENCODING", "GBK")


def run():
    driver = ogr.GetDriverByName('ESRI Shapefile')
    fileName = "../数据/数据.shp"  # 待拆分的shp文件
    dataSource = driver.Open(fileName, 0)
    layer = dataSource.GetLayer(0)
    layerDefinition = layer.GetLayerDefn()

    NameList = []
    TypeList = []
    WidthList = []
    for id in range(layerDefinition.GetFieldCount()):
        """获取属性表的列名/类型/长度"""
        NameList.append(layerDefinition.GetFieldDefn(id).GetName())
        TypeList.append(layerDefinition.GetFieldDefn(id).GetFieldTypeName(layerDefinition.GetFieldDefn(id).GetType()))
        WidthList.append(layerDefinition.GetFieldDefn(id).GetWidth())

    for i in range(0, layer.GetFeatureCount()):
        feat = layer.GetFeature(i)
        wkt = feat.geometry()
        name = feat.GetField('TBBH')  # 根据TBBH属性命名拆分后的文件

        crs = layer.GetSpatialRef()
        epsg = crs.GetAttrValue('AUTHORITY', 1)  # 获取shp文件的EPSG投影信息

        savepath = os.path.join('result', name)
        if not os.path.exists(savepath):
            os.mkdir(savepath)

        create_shp(name, str(wkt), epsg, feat, NameList, TypeList, WidthList, savepath)


def create_shp(name, wkt, epsg, feat, NameList, TypeList, WidthList, savepath):
    driver = ogr.GetDriverByName("ESRI Shapefile")
    data_source = driver.CreateDataSource(os.path.join(savepath, "{}.shp").format(name))
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(int(epsg))
    layer = data_source.CreateLayer(f'{name}', srs, ogr.wkbMultiPolygon)

    for id, name in enumerate(NameList):
        typ = eval('ogr.OFT' + TypeList[id])
        field_name = ogr.FieldDefn(name, typ)
        field_name.SetWidth(WidthList[id])
        layer.CreateField(field_name)
    feature = ogr.Feature(layer.GetLayerDefn())
    for id, name in enumerate(NameList):
        feature.SetField(name, feat.GetField(name))
    polygon = ogr.CreateGeometryFromWkt(wkt)
    feature.SetGeometry(polygon)
    layer.CreateFeature(feature)
    feature = None
    data_source = None


if __name__ == '__main__':
    run()
